from typing import List

from . import DEFAULT_MODEL_NAME, Prompt, client

SYSTEM_PROMPT = """你是一个有用的助手，负责将用户查询分类到预定义的意图类别中。
你的任务是确定用户查询是否匹配候选意图之一，代表其他意图，或者没有明确的意图。

对于每个用户查询，你应该用以下分类之一回应：
1. 候选意图之一（如果查询明确匹配该意图）
2. "other"（如果查询有明确意图但不匹配任何候选项）
3. "empty"（如果查询没有明确意图或太模糊）

只需回复分类标签，不要添加其他内容。"""

USER_PROMPT = """我需要你将以下用户查询分类为候选意图之一{other_option}{empty_option}。

候选意图：
{candidates}

用户查询：{user_input}

{examples}

分类："""

# 定义各种示例
EXAMPLES_ALL = """示例：
用户查询："我想了解一下个人所得税申报流程"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：other

用户查询："我想查询我的养老保险缴费记录"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：养老保险

用户查询："嗯..."
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：empty

用户查询："我在工作中受伤了，想申请工伤赔偿"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：工伤保险

用户查询："..."
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：empty"""

EXAMPLES_NO_EMPTY = """示例：
用户查询："我想了解一下个人所得税申报流程"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：other

用户查询："我想查询我的养老保险缴费记录"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：养老保险

用户查询："我在工作中受伤了，想申请工伤赔偿"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：工伤保险"""

EXAMPLES_NO_OTHER = """示例：
用户查询："我想查询我的养老保险缴费记录"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：养老保险

用户查询："嗯..."
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：empty

用户查询："我在工作中受伤了，想申请工伤赔偿"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：工伤保险

用户查询："..."
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：empty"""

EXAMPLES_ONLY_CANDIDATES = """示例：
用户查询："我想查询我的养老保险缴费记录"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：养老保险

用户查询："我在工作中受伤了，想申请工伤赔偿"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：工伤保险

用户查询："我被公司辞退了，想申请失业保险金"
候选意图：["养老保险", "失业保险", "工伤保险"]
分类：失业保险"""


def build_prompt(
    user_input: str,
    candidates: List[str],
    use_empty: bool = True,
    use_other: bool = True,
    use_examples: bool = True,
) -> Prompt:
    """Build the prompt for the triage agent.

    Args:
        user_input: The user's input query
        candidates: List of candidate intents
        use_empty: Whether to allow classification as "empty" intent
        use_other: Whether to allow classification as "other" intent
        use_examples: Whether to use examples

    Returns:
        Prompt object with system and user prompts
    """
    formatted_candidates = ", ".join([f'"{c}"' for c in candidates])

    other_option = '、"other"' if use_other else ""
    empty_option = '或"empty"' if use_empty else ""

    # 选择合适的示例
    if not use_examples:
        examples = ""
    elif use_empty and use_other:
        examples = EXAMPLES_ALL
    elif use_empty and not use_other:
        examples = EXAMPLES_NO_OTHER
    elif not use_empty and use_other:
        examples = EXAMPLES_NO_EMPTY
    else:
        examples = EXAMPLES_ONLY_CANDIDATES

    user_prompt = USER_PROMPT.format(
        user_input=user_input,
        candidates=formatted_candidates,
        other_option=other_option,
        empty_option=empty_option,
        examples=examples,
    )

    return Prompt(system_prompt=SYSTEM_PROMPT, user_prompt=user_prompt)


# 使用编辑距离找到最相似的候选意图
def levenshtein_distance(s1: str, s2: str) -> int:
    """计算两个字符串之间的编辑距离"""
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)

    if len(s2) == 0:
        return len(s1)

    previous_row = list(range(len(s2) + 1))
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            substitutions = previous_row[j] + (c1 != c2)
            current_row.append(min(insertions, deletions, substitutions))
        previous_row = current_row

    return previous_row[-1]


async def run_agent(
    user_input: str,
    candidates: List[str],
    use_empty: bool = True,
    use_other: bool = True,
    use_examples: bool = True,
) -> str:
    """Run the agent.

    Args:
        user_input: The user's input query
        candidates: List of candidate intents
        use_empty: Whether to allow classification as "empty" intent
        use_other: Whether to allow classification as "other" intent

    Returns:
        The classified intent as a string
    """
    prompt = build_prompt(user_input, candidates, use_empty, use_other, use_examples)

    chat_completion = await client.chat.completions.create(
        messages=[
            {"role": "system", "content": prompt.system_prompt},
            {"role": "user", "content": prompt.user_prompt},
        ],
        model=DEFAULT_MODEL_NAME,
        temperature=0,
    )

    response_str = chat_completion.choices[0].message.content

    # 使用编辑距离计算最相似的类别
    response_str = response_str.strip()

    # 如果响应正好是候选意图之一，直接返回
    if response_str in candidates:
        return response_str

    # 如果响应是 "other" 或 "empty"，直接返回
    if response_str in ["other", "empty"]:
        return response_str

    # 找到编辑距离最小的候选意图
    min_distance = float("inf")
    closest_intent = response_str

    for candidate in candidates:
        distance = levenshtein_distance(response_str, candidate)
        if distance < min_distance:
            min_distance = distance
            closest_intent = candidate

    # 如果最小编辑距离小于阈值（比如2），返回最相似的候选
    if min_distance <= 2 or closest_intent in response_str:
        return closest_intent

    # 否则，如果允许 "other"，返回 "other"
    if use_other:
        return "other"

    # 如果不允许 "other"，返回第一个候选作为默认值
    return candidates[0] if candidates else "empty"
